package vip.kqsy.cloud.ai.server.websocket;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.MessageManager;
import com.alibaba.dashscope.common.Role;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;
import vip.kqsy.cloud.ai.constant.em.AiSessionHistoryTypeEm;

import vip.kqsy.cloud.ai.server.pojo.entity.AiSession;
import vip.kqsy.cloud.ai.server.pojo.entity.AiSessionHistory;
import vip.kqsy.cloud.ai.server.service.AiSessionHistoryService;
import vip.kqsy.cloud.ai.server.service.AiSessionService;
import vip.kqsy.cloud.ai.server.utils.HtmlUtil;
import vip.kqsy.cloud.core.common.context.ApplicationContextHolder;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Comparator;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

@Component
@Slf4j
@ServerEndpoint("/ws/chat/{sessionId}")
public class ChatWs  {

    private AiSessionService aiSessionService;
    private AiSessionHistoryService aiSessionHistoryService;
    private Generation gptService;

    private Long sessionId;
    private Session session;

    @OnOpen
    public void onOpen(Session session, @PathParam("sessionId") Long sessionId) {
        log.info("WS对话：{}", sessionId);
        this.sessionId = sessionId;
        this.session = session;
        aiSessionService = ApplicationContextHolder.get().getBean(AiSessionService.class);
        aiSessionHistoryService = ApplicationContextHolder.get().getBean(AiSessionHistoryService.class);
        gptService = ApplicationContextHolder.get().getBean(Generation.class);
    }

    @OnMessage
    @SneakyThrows
    public void onMessage(String message) throws IOException {
        log.info("收到消息--{}", message);
        AiSession aiSession = aiSessionService.getById(sessionId);
        Assert.notNull(aiSession, "上下文不存在");
        AiSessionHistory history = new AiSessionHistory();
        history.setContent(message);
        history.setSessionId(sessionId);
        history.setType(AiSessionHistoryTypeEm.USER.getCode());
        history.setCreateTime(LocalDateTime.now());
        aiSessionHistoryService.save(history);

//        CreateChatReq req = new CreateChatReq();
//        Message message1 = new Message();
//        message1.setRole(Message.ROLE_USER);
//        message1.setContent(message);
//        req.setMessages(Collections.singletonList(message1));

        GenerationResult chatRes = gptService.call(build(message));
        String resText = chatRes.getOutput().getText();
        AiSessionHistory history2 = new AiSessionHistory();
        history2.setSessionId(sessionId);
        history2.setType(AiSessionHistoryTypeEm.SYSTEM.getCode());
        history2.setCreateTime(LocalDateTime.now());
        history2.setContent(resText);
        aiSessionHistoryService.save(history2);

        session.getBasicRemote().sendText(HtmlUtil.convert(history2.getContent()));
    }

    @OnClose
    public void onClose() {
        log.info("WS结束：{}", sessionId);
    }


    private GenerationParam build(String text) {

        MessageManager manager = new MessageManager();

        List<AiSessionHistory> list = aiSessionHistoryService.lambdaQuery()
                .eq(AiSessionHistory::getSessionId, sessionId)
                .orderByDesc(AiSessionHistory::getCreateTime)
                .list();
        list.sort(Comparator.comparing(AiSessionHistory::getCreateTime));

        list.forEach(item -> {
            Message message = new Message();
            if (Objects.equals(AiSessionHistoryTypeEm.SYSTEM.getCode(), item.getType())) {
                message.setRole(Role.SYSTEM.getValue());
            } else {
                message.setRole(Role.USER.getValue());
            }
            message.setContent(item.getContent());
            manager.add(message);
        });

        // 加上最新的用户
        Message message1 = new Message();
        message1.setRole(Role.USER.getValue());
        message1.setContent(text);
        manager.add(message1);

        return GenerationParam.builder()
                .model("qwen1.5-72b-chat")
                .messages(manager.get())
                .build();
    }

}
